from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict
import requests

class PaperRetriever:
    """论文检索器"""
    def __init__(self, base_url: str = "http://180.184.65.98:38880/atomgit"):
        self.base_url = base_url
        self.executor = ThreadPoolExecutor(max_workers=4)

    async def batch_retrieve(self, queries: List[str], top_k: int = 30) -> List[Dict]:
        """并行检索多个查询"""
        futures = []
        for query in queries:
            futures.append(
                self.executor.submit(self.search_papers, query, top_k)
            )
        return [f.result() for f in futures]

    def search_papers(self, query: str, top_k: int = 30) -> List[Dict]:
        """检索论文"""
        url = f"{self.base_url}/search_papers"
        params = {"query": query, "top_k": top_k}
        try:
            response = requests.get(url, params=params)
            if response.status_code == 200:
                return response.json()
        except Exception as e:
            print(f"检索失败: {str(e)}")
        return []

    def expand_citations(self, text: str) -> List[Dict]:
        """扩展引用来源"""
        # 检索标题中提到的相关论文
        title_papers = self.search_by_title_contain(text)
        # 检索内容中提到的相关段落
        chunk_papers = self.search_by_chunk_contain(text)
        # 使用title和chunk_id的组合作为唯一键，保留同一篇论文不同chunk的引用
        return list({f"{p['title']}_{p.get('chunk_id', '')}": p for p in title_papers + chunk_papers}.values())

    def search_by_title_contain(self, text: str, top_k: int = 30) -> List[Dict]:
        """检索标题中包含特定文本的论文"""
        url = f"{self.base_url}/query_paper_metadata_that_title_contain"
        params = {"title": text, "top_k": top_k}
        try:
            response = requests.get(url, params=params)
            if response.status_code == 200:
                return response.json()
        except Exception as e:
            print(f"检索标题失败: {str(e)}")
        return []

    def search_by_chunk_contain(self, text: str, top_k: int = 30) -> List[Dict]:
        """检索内容中包含特定文本的段落"""
        url = f"{self.base_url}/search_papers"
        params = {"query": text, "top_k": top_k}
        try:
            response = requests.get(url, params=params)
            if response.status_code == 200:
                return response.json()
        except Exception as e:
            print(f"检索段落失败: {str(e)}")
        return []
